from inspect import isfunction

import torch
import torch.nn as nn
import torch.nn.functional as F


class Identity(nn.Module):
    """
    Identity block.
    """
    def __init__(self):
        super(Identity, self).__init__()

    def forward(self, x):
        return x


class Swish(nn.Module):
    """
    Swish activation function from 'Searching for Activation Functions.
    https://arxiv.org/abs/1710.05941.
    """
    def forward(self, x):
        return x * torch.sigmoid(x)


class HSigmoid(nn.Module):
    """
    Approximated sigmoid function,
        so-called hard-version of sigmoid from 'Searching for MobileNetV3'.
    https://arxiv.org/abs/1905.02244.
    """
    def forward(self, x):
        return F.relu6(x + 3.0, inplace=True) / 6.0


class HSwish(nn.Module):
    """
    H-Swish activation function from 'Searching for MobileNetV3'.
    https://arxiv.org/abs/1905.02244.

    Args:
        inplace (bool, optional):
            Whether to use inplace version of the module. Defaults to False.
    """
    def __init__(self, inplace=False):
        super(HSwish, self).__init__()
        self.inplace = inplace

    def forward(self, x):
        return x * F.relu6(x + 3.0, inplace=self.inplace) / 6.0


def get_activation_layer(activation):
    assert (activation is not None)
    if isfunction(activation):
        return activation()
    elif isinstance(activation, str):
        if activation == "relu":
            return nn.ReLU(inplace=True)
        elif activation == "relu6":
            return nn.ReLU6(inplace=True)
        elif activation == "swish":
            return Swish()
        elif activation == "hswish":
            return HSwish(inplace=True)
        elif activation == "sigmoid":
            return nn.Sigmoid()
        elif activation == "hsigmoid":
            return HSigmoid()
        elif activation == "identity":
            return Identity()
        else:
            raise NotImplementedError()
    else:
        assert (isinstance(activation, nn.Module))
        return activation


class ConvBlock(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride,
                 padding,
                 dilation=1,
                 groups=1,
                 bias=False,
                 use_bn=True,
                 bn_eps=1e-5,
                 activation=(lambda: nn.ReLU(inplace=True))):
        """Standard convolution block with Batch normalization and activation.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            kernel_size ():
                Convolution window size.
            stride (int | List[int, int] | Tuple[int, int]):
                Strides of the convolution.
            padding (int | List/Tuple of 2 int | List/Tuple of 4 int):
                Padding value for convolution layer.
            dilation (int | List[int, int] | Tuple[int, int]):
                Dilation value for convolution layer. Defaults to 1.
            groups (int, optional):
                Number of groups. Defaults to 1.
            bias (bool, optional):
                Whether the layer uses a bias vector. Defaults to False.
            use_bn (bool, optional):
                Whether to use BatchNorm layer. Defaults to True.
            bn_eps (float, optional):
                Small float added to variance in Batch norm. Defaults to 1e-5.
            activation (function | str | None, optional):
                Activation function or name of activation function.
                Defaults to (lambda: nn.ReLU(inplace=True)).
        """
        super(ConvBlock, self).__init__()
        self.activate = (activation is not None)
        self.use_bn = use_bn
        self.use_pad = (isinstance(padding, (list, tuple))
                        and (len(padding) == 4))

        if self.use_pad:
            self.pad = nn.ZeroPad2d(padding=padding)
            padding = 0
        self.conv = nn.Conv2d(in_channels=in_channels,
                              out_channels=out_channels,
                              kernel_size=kernel_size,
                              stride=stride,
                              padding=padding,
                              dilation=dilation,
                              groups=groups,
                              bias=bias)
        if self.use_bn:
            self.bn = nn.BatchNorm2d(num_features=out_channels, eps=bn_eps)
        if self.activate:
            self.activ = get_activation_layer(activation)

    def forward(self, x):
        if self.use_pad:
            x = self.pad(x)
        x = self.conv(x)
        if self.use_bn:
            x = self.bn(x)
        if self.activate:
            x = self.activ(x)
        return x


def conv1x1_block(in_channels,
                  out_channels,
                  stride=1,
                  padding=0,
                  groups=1,
                  bias=False,
                  use_bn=True,
                  bn_eps=1e-5,
                  activation=(lambda: nn.ReLU(inplace=True))):
    return ConvBlock(in_channels=in_channels,
                     out_channels=out_channels,
                     kernel_size=1,
                     stride=stride,
                     padding=padding,
                     groups=groups,
                     bias=bias,
                     use_bn=use_bn,
                     bn_eps=bn_eps,
                     activation=activation)


def conv3x3_block(in_channels,
                  out_channels,
                  stride=1,
                  padding=1,
                  dilation=1,
                  groups=1,
                  bias=False,
                  use_bn=True,
                  bn_eps=1e-5,
                  activation=(lambda: nn.ReLU(inplace=True))):
    return ConvBlock(in_channels=in_channels,
                     out_channels=out_channels,
                     kernel_size=3,
                     stride=stride,
                     padding=padding,
                     dilation=dilation,
                     groups=groups,
                     bias=bias,
                     use_bn=use_bn,
                     bn_eps=bn_eps,
                     activation=activation)


def conv7x7_block(in_channels,
                  out_channels,
                  stride=1,
                  padding=3,
                  bias=False,
                  use_bn=True,
                  activation=(lambda: nn.ReLU(inplace=True))):
    return ConvBlock(in_channels=in_channels,
                     out_channels=out_channels,
                     kernel_size=7,
                     stride=stride,
                     padding=padding,
                     bias=bias,
                     use_bn=use_bn,
                     activation=activation)


class ResBlock(nn.Module):
    """
    Simple ResNet block for residual path in ResNet unit.

    Args:
        in_channels (int):
            Number of input channels.
        out_channels (int):
            Number of output channels.
        stride (int or tuple/list of 2 int):
            Strides of the convolution.
        bias (bool, default False):
            Whether the layer uses a bias vector.
        use_bn (bool, default True):
            Whether to use BatchNorm layer.

    """
    def __init__(self,
                 in_channels,
                 out_channels,
                 stride,
                 bias=False,
                 use_bn=True):
        super(ResBlock, self).__init__()
        self.conv1 = conv3x3_block(
            in_channels=in_channels,
            out_channels=out_channels,
            stride=stride,
            bias=bias,
            use_bn=use_bn)
        self.conv2 = conv3x3_block(
            in_channels=out_channels,
            out_channels=out_channels,
            bias=bias,
            use_bn=use_bn,
            activation=None)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x
